jax
tensorflow
flax
tqdm
numpy